# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
from collections import OrderedDict

import tabulate
from tornado.ioloop import PeriodicCallback


def iter_dict(dict_obj, sort_keys):
  if not sort_keys:
    for k, v in dict_obj.items():
      yield k, v
  else:
    keys = sorted(dict_obj.keys())
    for k in keys:
      yield k, dict_obj[k]


class PeriodicStatusPrinter(object):
  def __init__(self, sort_keys=False, logger=None):
    self._logger = logger or logging.getLogger('PeriodicStatusPrinter')
    self._periodic_callback = None
    self._status_update_callback = []
    self._status = OrderedDict()
    self._sort_keys = sort_keys

  def add_status_update_callback(self, callback):
    self._status_update_callback.append(callback)

  @property
  def running(self):
    return self._periodic_callback is not None

  def start(self, period_ms=10000):
    if self.running:
      return
    self._periodic_callback = PeriodicCallback(self.print_status, period_ms)
    self._periodic_callback.start()

  def stop(self):
    if not self.running:
      return
    self._periodic_callback.stop()
    self._periodic_callback = None

  def print_status(self, sort_keys=None, ignore_error=False):
    try:
      update_time = datetime.datetime.now()
      for callback in self._status_update_callback:
        callback(update_time, self._status)
      self._print_status(update_time, sort_keys=sort_keys)
    except Exception:
      if not ignore_error:
        raise
      self._logger.exception('Exception raised')

  def _print_status(self, update_time, sort_keys=None):
    update_time = update_time or datetime.datetime.now()
    sort_keys = self._sort_keys if sort_keys is None else sort_keys
    table = []

    for key, value in iter_dict(self._status, sort_keys):
      table.append(['+', key, value])

    status_table_str = tabulate.tabulate(table, headers=['+', 'Key', 'Value'])
    self._logger.info('----- Status:\n+++ STATUS: %s\n%s\n+++', str(update_time), status_table_str)
